#!/usr/bin/env python3
import argparse
import csv
import os
import sys
import shutil
from Bio import SeqIO
from Bio.SeqIO.FastaIO import FastaIterator

def main():
    parser = argparse.ArgumentParser(description='Renumber FASTA headers based on taxonomy')
    parser.add_argument('-t', required=True, help='Taxonomy TSV file')
    parser.add_argument('-f', required=True, help='Input FASTA file')
    parser.add_argument('-o', help='Output CSV file (default: <input_file>.header_changes.csv)')
    parser.add_argument('-a', action='append', help='Annotation fields from Taxonomy file to include in header', default=[])
    args = parser.parse_args()

    # 设置默认输出文件名
    if not args.o:
        args.o = f"{args.f}.header_changes.csv"

    # 验证输入文件
    if not all([os.path.exists(args.t), os.path.exists(args.f)]):
        sys.exit("错误：输入文件不存在")

    # 读取Taxonomy数据
    taxonomy = {}
    with open(args.t, 'r') as tsv:
        reader = csv.DictReader(tsv, delimiter='\t')
        # 验证注释字段是否存在
        if args.a:
            fieldnames = reader.fieldnames
            for field in args.a:
                if field not in fieldnames:
                    sys.exit(f"错误：Taxonomy文件中不存在'{field}'字段")
        
        for row in reader:
            taxonomy[row['Prefix']] = row

    # 获取文件名前缀匹配
    base_name = os.path.basename(args.f)
    file_prefix = base_name.split('.')[0]
    
    if file_prefix not in taxonomy:
        sys.exit(f"错误：未找到匹配的Prefix '{file_prefix}'")

    abbr = taxonomy[file_prefix]['ABBR']
    counter = 1

    # 处理FASTA并生成CSV
    # 创建备份文件
    shutil.copy2(args.f, f"{args.f}.original")
    
    # 修改输出文件为原文件名
    with open(args.f, 'w') as fasta_out, \
         open(args.o, 'w') as csv_out:
    
        writer = csv.writer(csv_out)
        writer.writerow(['Original Header', 'New Header'])
    
        # 从备份文件读取数据
        with open(f"{args.f}.original") as in_fasta:
            for record in SeqIO.parse(in_fasta, 'fasta'):
                # 保存原始header
                original_header = record.description if record.description else record.id
                
                # 生成基本ID
                new_id = f"{abbr}{counter:06d}"
                
                # 添加注释信息
                if args.a:
                    annotations = [taxonomy[file_prefix][field] for field in args.a]
                    full_header = f"{new_id} | {' | '.join(annotations)}"
                    # 保存完整描述，但ID仍然保持纯数字格式
                    record.id = new_id
                    record.description = f"| {' | '.join(annotations)}"
                else:
                    full_header = new_id
                    record.id = new_id
                    record.description = ''
                
                # 写入CSV时使用原始header和新header
                writer.writerow([original_header, full_header])
                
                # 写入新数据到原文件
                SeqIO.write(record, fasta_out, 'fasta')
                
                counter += 1

    print(f"处理完成！生成文件：{args.f}已重新编号，同时生成了备份{args.o} 和 {args.f}.original")

if __name__ == '__main__':
    main()